import cv2
import numpy as np
import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from moviepy.editor import VideoFileClip
%matplotlib inline
def undistort(image, gray):
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objectpoints, imagepoints, gray.shape[::-1], None, None)
return cv2.undistort(image, mtx, dist, None, mtx)
for i in range(1, 2):
fname = test_images_dir + "/test" + str(i) + ".jpg"
image = mpimg.imread(fname)
gray = grayscale(image)
undist = undistort(image, gray)
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 9))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original', fontsize=50)
ax2.imshow(undist)
ax2.set_title('Undistorted', fontsize=50)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)
def calibrate_chessboard(x_dim=9, y_dim=6, drawCorners=0, writeToFile=0):
out = [] # Output images
objp = np.zeros((y_dim * x_dim, 3), np.float32)
objp[:, :2] = np.mgrid[0:x_dim, 0:y_dim].T.reshape(-1, 2) # x, y coordinates
images = glob.glob(calibration_dir + '/calibration*.jpg')
for fname in images:
img = mpimg.imread(fname)
gray = grayscale(img)
ret, corners = cv2.findChessboardCorners(gray, (x_dim, y_dim), None)
if ret:
imagepoints.append(corners)
objectpoints.append(objp)
if drawCorners:
img = cv2.drawChessboardCorners(img, (x_dim, y_dim), corners, ret)
undist = undistort(img, gray)
out.append(undist)
if writeToFile:
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 9))
f.tight_layout()
ax1.imshow(img)
ax1.set_title('Original', fontsize=50)
ax2.imshow(undist)
ax2.set_title('Undistorted', fontsize=50)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)
#plt.savefig(output_dir + '/out_' + fname.split('/')[-1])
return out
out = calibrate_chessboard(drawCorners=1, writeToFile=1)